package com.btm.mcpserver;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

@Component
public class DatabaseMcpServer {
    private static final Logger log = LoggerFactory.getLogger(DatabaseMcpServer.class);

    @Autowired
    private JdbcTemplate jdbcTemplate;
    
    @Autowired
    private ObjectMapper objectMapper;

    /**
     * Tool注解中的description的内容，决定着大模型如何使用这个mcp-server
     * @param querySQL 要执行的SQL
     * @return 查询结果的JSON字符串表示
     */
    @Tool(description = "执行数据库查询，返回JSON格式结果")
    public String executeQuerySql(String querySQL) {
        log.info("收到待执行SQL：\n{}", querySQL);
        if (querySQL.toLowerCase().trim().startsWith("select")) {
            try {
                List<Map<String, Object>> results = jdbcTemplate.queryForList(querySQL);
                if (results.isEmpty()) {
                    return "{\"message\": \"查询结果为空\"}";
                }

                // 打印查询结果到info日志
        //                log.info("查询结果: {}", results);
                log.info("查询结果数据量：{}", results.size());
                try {
                    // 将查询结果直接转换为JSON字符串
                    return objectMapper.writeValueAsString(results);
                } catch (JsonProcessingException e) {
                    log.error("转换JSON出错: {}", e.getMessage(), e);
                    return "{\"error\": \"转换JSON出错: " + e.getMessage() + "\"}";
                }
            } catch (Exception e) {
                log.error("执行SQL出错: {}", e.getMessage(), e);
                return "{\"error\": \"执行SQL出错: " + e.getMessage() + "\"}";
            }
        } else {
            log.error("非SELECT语句，暂不支持执行");
            return "{\"error\": \"非SELECT语句，暂不支持执行\"}";
        }
    }

    @Tool(description = "执行所有类型的SQL")
    public String executeAllSql(String sql) {
        try {
            String sqlType = sql.trim().toLowerCase().split("\\s+")[0];
            if ("select".equals(sqlType)) {
                // 处理 SELECT 查询
                List<Map<String, Object>> results = jdbcTemplate.queryForList(sql);
                if (results.isEmpty()) {
                    return "{\"message\": \"查询结果为空\"}";
                }
                log.info("查询结果数据量：{}", results.size());
                return objectMapper.writeValueAsString(results);
            } else {
                // 处理非 SELECT 语句（如 INSERT、UPDATE、DELETE 等）
                int affectedRows = jdbcTemplate.update(sql);
                return "{\"message\": \"执行成功\", \"affectedRows\": " + affectedRows + "}";
            }
        } catch (Exception e) {
            log.error("执行SQL出错: {}", e.getMessage(), e);
            return "{\"error\": \"执行SQL出错: " + e.getMessage() + "\"}";
        }
    }
}
